/**
 * 该模块实现了一个红黑树。
 * License: BSD
 * Authors: Lucifer (786325481@QQ.com)
 * Copyright: Copyright (C) 2008 Lucifer. All rights reserved.
 */

/**
 * 更多详情请参考 www.cs.princeton.edu/~rs/talks/LLRB/
 */

module system.collections.RBTree;

import system.Traits;
import system.Exception;

import system.collections.IComparer;
import system.collections.Comparer;

version(unittest)
{
	import system.collections.Stack;
}

package enum Color : bool
{
	Red = true,
	Black = false
}

/**
 * 红黑树节点
 */
public final class RBTreeNode(T)
{
	private T theValue;
	private Color theColor;
	private RBTreeNode!(T) theLeft, theRight;

	this(T value)
	{
		theValue = value;
		theColor = Color.Red;
	}

	/**
	 * 返回此节点的左子树。
	 */
	public RBTreeNode!(T) left()
	{
		return theLeft;
	}

	/**
	 * 返回此节点的右子树。
	 */
	public RBTreeNode!(T) right()
	{
		return theRight;
	}

	/**
	 * 返回此节点的值。
	 */
	public T value()
	{
		return theValue;
	}

	/**
	 * 设置此节点的值。
	 */
	public void value(T newValue)
	{
		theValue = newValue;
	}
}

public final class RBTree(T)
{
    static if( is( T == struct ) && !( is ( typeof(T.opCmp) ) ) )
    {
        static assert(false, "Struct " ~ T.stringof ~ " should overload opCmp");
    }

	private RBTreeNode!(T) theRoot;		// root of the BST
	private IComparer!(T) theComparer;
	private int theSize;

	public this(IComparer!(T) comparer = DefaultComparer!(T).getInstance())
	in
	{
		assert(comparer !is null);
	}
	body
	{
		theComparer = comparer;
	}

	/**
	 * 检测红黑是否包含 item 。
	 * Params:
	 * 		item = 要检测的项。
	 * Returns:
	 *		如果包含，返回 true；否则，返回 false 。
	 */
	public bool contains(T item)
	{
		return findNode(item) !is null;
	}

	/**
	 * 依据指定项查找节点。
	 * Params:
	 * 		item = 要查找的项。
	 * Returns:
	 *		如果找到，返回该红黑树节点；否则，返回 null 。
	 */
	public RBTreeNode!(T) findNode(T item)
	{
		int result;

		for(auto node = theRoot; node !is null; node = (result < 0) ? node.theLeft : node.theRight)
		{
			result = theComparer.compare(item, node.value);
			if(result == 0)
				return node;
		}

		return null;
	}

	/**
	 * 返回此红黑树中最小的节点。
	 * Returns:
	 *		如果红黑树为空，返回 null；否则，返回最小节点。
	 */
	public RBTreeNode!(T) min()
	{
		if(theRoot is null)
			return null;
		return min(theRoot);
	}

	/**
	 * 返回此红黑树中最大的节点。
	 * Returns:
	 *		如果红黑树为空，返回 null；否则，返回最大节点。
	 */
	public RBTreeNode!(T) max()
	{
		if(theRoot is null)
			return null;
		return max(theRoot);
	}

	/**
	 * 返回此红黑树的根。
	 * Returns:
	 *		如果红黑树为空，返回 null；否则，返回根。
	 */
	public RBTreeNode!(T) root()
	{
		return theRoot;
	}

	/**
	 * 返回此红黑树所包含的节点数。
	 */
	public int count()
	{
		return theSize;
	}

	/**
	 * 添加指定的项到红黑树中。
	 * Params:
	 *		item = 要添加的项。
	 * Throws:
	 *		ArgumentException，当有重复值插入时。
	 * Comments:
	 *		当没有正确实现 opCmp 时，会抛出比较错误。
	 */
	public void add(T item)
	{
		theRoot = insert(theRoot, item);
		theRoot.theColor = Color.Black;
	}

	/**
	 * 从此红黑树中移除最小的节点。
	 */
	public void removeMin()
	{
		if(theRoot !is null)
		{
			theRoot = removeMin(theRoot);
			if(theRoot !is null)
				theRoot.theColor = Color.Black;
		}
	}

	/**
	 * 从此红黑树中移除最大的节点。
	 */
	public void removeMax()
	{
		if(theRoot !is null)
		{
			theRoot = removeMax(theRoot);
			if(theRoot !is null)
				theRoot.theColor = Color.Black;
		}
	}

	/**
	 * 从此红黑树中移除指定的项。
	 * Params:
	 * 		item = 要移除的项。
	 */
	public void remove(T item)
	{
		if(theRoot !is null)
		{
			theRoot = remove(theRoot, item);
			if(theRoot !is null)
				theRoot.theColor = Color.Black;
		}
	}

	/**
	 * 清空此红黑树。
	 */
	public void clear()
	{
		/*
		 * 我们可以应用 D 语言关键字 delete 重新实现清空操作。
		 * 但是通过 GC 来回收或许是更好的解决方案。
		 */
		theRoot = null;
		theSize = 0;
	}

	/**
	 * 返回此红黑树所使用的比较器。
	 */
	public IComparer!(T) comparer()
	{
		return theComparer;
	}

	//检测节点的颜色是否为红色
	private bool isRed(RBTreeNode!(T) node)
	{
		if(node is null)
			return false;
		return node.theColor == Color.Red;
	}

	/*
	 * 插入指定的项到指定的节点，并返回该节点。
	 */
	private RBTreeNode!(T) insert(RBTreeNode!(T) node, T item)
	{
		// insert a new node at the bottom.
		if(node is null)
		{
			theSize++;
			return new RBTreeNode!(T)(item);
		}

		// split 4-nodes on the way down
		if(isRed(node.theLeft) && isRed(node.theRight))
			colorFlip(node);

		// standard BST insert code
		auto result = theComparer.compare(item, node.value);

		if(result == 0)
			node.value = item;
			//throw new ArgumentException("The value adding is duplicate.");
		else if(result < 0)
			node.theLeft = insert(node.theLeft, item);
		else
			node.theRight = insert(node.theRight, item);

		// fix right-leaning reds on the way up
		if(isRed(node.theRight))
			node = rotateLeft(node);

		// fix two reds in a row on the way up
		if(isRed(node.theLeft) && isRed(node.theLeft.theLeft))
			node = rotateRight(node);

		return node;
	}

	private RBTreeNode!(T) rotateLeft(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		// Make a right-leaning 3-node lean to the left.
		auto RBNode = node.theRight;
		node.theRight = RBNode.theLeft;
		RBNode.theLeft = node;
		RBNode.theColor = RBNode.theLeft.theColor;
		RBNode.theLeft.theColor = Color.Red;
		return RBNode;
	}

	private RBTreeNode!(T) rotateRight(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		// Make a left-leaning 3-node lean to the right.
		auto RBNode = node.theLeft;
		node.theLeft = RBNode.theRight;
		RBNode.theRight = node;
		RBNode.theColor = RBNode.theRight.theColor;
		RBNode.theRight.theColor = Color.Red;
		return RBNode;
	}

	// Flip the colors of the three nodes
	private void colorFlip(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		node.theColor = cast(Color)(!node.theColor);
		node.theLeft.theColor = cast(Color)(!node.theLeft.theColor);
		node.theRight.theColor = cast(Color)(!node.theRight.theColor);
	}

	// Fix right-leaning reds and eliminate 4-nodes on the way up.
	private RBTreeNode!(T) fixUp(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		// rotate-left right-leaning reds
		if(isRed(node.theRight))
			node = rotateLeft(node);

		// rotate-right red-red pairs
		if(isRed(node.theLeft) && isRed(node.theLeft.theLeft))
			node = rotateRight(node);

		// split 4-nodes
		if(isRed(node.theLeft) && isRed(node.theRight))
			colorFlip(node);

		return node;
	}

	private RBTreeNode!(T) moveRedRight(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		// Assuming that node is red and both node.right and node.right.left
		// are black, make node.right or one of its children red.
		colorFlip(node);
		if(isRed(node.theLeft.theLeft))
		{
			node = rotateRight(node);
			colorFlip(node);
		}

		return node;
	}

	private RBTreeNode!(T) moveRedLeft(RBTreeNode!(T) node)
	in
	{
		assert(node !is null);
	}
	body
	{
		// Assuming that node is red and both node.left and node.left.left
		// are black, make node.left or one of its children red.
		colorFlip(node);
		if(isRed(node.theRight.theLeft))
		{
			node.theRight = rotateRight(node.theRight);
			node = rotateLeft(node);
			colorFlip(node);
		}
		return node;
	}

	private RBTreeNode!(T) removeMax(RBTreeNode!(T) node)
	{
		// lean 3-nodes to the right
		if(isRed(node.theLeft))
			node = rotateRight(node);

		// remove node on bottom level
		// node must be Color.Red by invariant
		if(node.theRight is null)
		{
			theSize--;
			return null;
		}

		// borrow from sibling if necessary
		if(!isRed(node.theRight) && !isRed(node.theRight.theLeft))
			node = moveRedRight(node);

		// move down one level
		node.theRight = removeMax(node.theRight);

		// fix right-leaning red links
		// and eliminate 4-nodes
		// on the way up
		return fixUp(node);
	}

	private RBTreeNode!(T) removeMin(RBTreeNode!(T) node)
	{
		// remove node on bottom level
		// node must be Color.Red by invariant
		if(node.theLeft is null)
		{
			theSize--;
			return null;
		}

		// push red link down if necessary
		if(!isRed(node.theLeft) && !isRed(node.theLeft.theLeft))
			node = moveRedLeft(node);

		// move down one level
		node.theLeft = removeMin(node.theLeft);

		// fix right-leaning red links
		// and eliminate 4-nodes
		// on the way up
		return fixUp(node);
	}

	private RBTreeNode!(T) remove(RBTreeNode!(T) node, T item)
	{
		if(node is null)
			return null;

		auto result = theComparer.compare(item, node.value);
		if(result < 0)
		{// LEFT
			// push red right if necessary
			auto left = node.theLeft;
			if(left !is null)
			{
				auto leftNode = left.theLeft;
				if(!isRed(left) && !isRed(leftNode))
					node = moveRedLeft(node);
			}
			// move down (left)
			node.theLeft = remove(node.theLeft, item);
		}
		else
		{// RIGHT or EQUAL
			// rotate to push red right
			if(isRed(node.theLeft))
				node = rotateRight(node);

			// EQUAL (at bottom)
			// delete node
			if((result == 0) && (node.theRight is null))
			{
				theSize--;
				return null;
			}

			// push red right if necessary
			auto right = node.theRight;
			if(right !is null)
			{
				auto leftNode = right.theLeft;
				if(!isRed(right) && !isRed(leftNode))
					node = moveRedRight(node);
			}

			// EQUAL (not at bottom)
			if(result == 0)
			{
				// replace current node with successor
				node.value = min(node.theRight).value;
				// delete successor
				node.theRight = removeMin(node.theRight);
			}
			else
				// move down (right)
				node.theRight = remove(node.theRight, item);
		}
		// fix right-leaning red links
		// and eliminate 4-nodes
		// on the way up
		return fixUp(node);
	}

	private RBTreeNode!(T) min(RBTreeNode!(T) node)
	in
	{
		assert( node !is null );
	}
	body
	{
		auto RBNode = node;
		while(RBNode.theLeft !is null)
		{
			RBNode = RBNode.theLeft;
		}
		return RBNode;
	}

	private RBTreeNode!(T) max(RBTreeNode!(T) node)
	in
	{
		assert( node !is null );
	}
	body
	{
		auto RBNode = node;
		while(RBNode.theRight !is null)
		{
			RBNode = RBNode.theRight;
		}
		return RBNode;
	}

	version(unittest)
	{
		private bool isRBTree()
		{
			return isBST() && is234Tree() && isBalanced();
		}

		private bool isBST()
		{
			return isBST(theRoot, min(), max());
		}

		private bool is234Tree()
		{
			return is234Tree(theRoot);
		}

		private bool isBalanced()
		{
			// Do all paths from root to leaf have same number of black edges?
			auto black = 0;	// number of black links on path from root to min
			auto node = theRoot;

			while( node !is null )
			{
				if(!isRed(node))
					black++;
				node = node.left;
			}

			return isBalanced(theRoot, black);
		}

		private bool isBST(RBTreeNode!(T) node, RBTreeNode!(T) minNode, RBTreeNode!(T) maxNode)
		{
			// Are all the values in the BST rooted at node between min and max,
			// and does the same property hold for both subtrees?

			if(node is null)
				return true;

			if(theComparer.compare(node.value, minNode.value) < 0 ||
				theComparer.compare(node.value, maxNode.value) > 0)
				return false;
			return isBST(node.left, minNode, node) && isBST(node.right, node, maxNode);
		}

		private bool is234Tree(RBTreeNode!(T) node)
		{
			// Does the tree have no red right links, and at most two (left)
			// red links in a row on any path?

			if(node is null)
				return true;

			if(isRed(node.right))
				return false;

			if(isRed(node))
				if(isRed(node.left))
					if(isRed(node.left.left))
						return false;

			return is234Tree(node.left) && is234Tree(node.right);
		}

		private bool isBalanced(RBTreeNode!(T) node, int black)
		{
			// Does every path from the root to a leaf have the given number
			// of black links?

			if(node is null && black == 0)
				return true;
			else if(node is null && black != 0)
				return false;

			if(!isRed(node))
				black--;

			return isBalanced(node.left, black) && isBalanced(node.right, black);
		}
	}
}

unittest
{
	auto tree = new RBTree!(int)();

	for(auto i = 1; i <= 100; i++)
	{
		tree.add(i);
	}

	auto stack = new Stack!(RBTreeNode!(int))(100);
	RBTreeNode!(int) current = null;

	for(auto node = tree.theRoot; node !is null; node = node.left)
	{
		stack.push(node);
	}

	bool moveNext()
	{
		if(stack.count == 0)
		{
			current = null;
			return false;
		}
		current = stack.pop();
		for(auto node = current.right; node !is null; node = node.left)
		{
			stack.push(node);
		}
		return true;
	}

	int currentValue()
	{
		if(current !is null)
			return current.value;
		return int.init;
	}

	int i = 1;
	while(moveNext())
	{
		assert(i++ == currentValue());
	}

	assert(100 == tree.max().value);
	assert(1 == tree.min().value);

	assert(tree.count == 100);

	tree.removeMax();

	assert(tree.count == 100-1);
	assert(tree.max.value == 99);
	assert(!tree.contains(100));

	tree.removeMin();

	assert(tree.count == 100-2);
	assert(tree.min.value == 2);
	assert(!tree.contains(1));

	tree.remove(2);
	tree.remove(15);
	tree.remove(22);

	assert(!tree.contains(2));
	assert(!tree.contains(15));
	assert(!tree.contains(22));

	assert(tree.count == 100-5);

	assert(tree.min.value == 3);
	assert(tree.max.value == 99);

	tree.remove(2);
	tree.remove(15);
	tree.remove(22);

	assert(tree.count == 100-5);

	assert(tree.isBST());
	assert(!tree.is234Tree());
	assert(tree.isBalanced());

	tree.clear();

	assert(tree.count == 0);

	tree.removeMin();
	tree.removeMax();

	tree.remove(1);
}
